import numpy as np
import pandas as pd
from scipy.linalg import eigh
import matplotlib.pyplot as plt

def generate_DOAs(num_doa):
    """
    生成指定数量的随机 DOA，每个 DOA 都是 10 的整数倍
    """
    # 设置 DOA 的范围，例如 -90 到 90 度
    min_doa = -90
    max_doa = 90
    
    # 生成范围内的所有 10 的整数倍
    possible_DOAs = np.arange(min_doa, max_doa + 1, 10)
    
    # 从这些可能的值中随机选择
    DOAs = np.random.choice(possible_DOAs, size=num_doa, replace=True)
    
    return DOAs

def steering_vector(M, angle, wavelength=1, d=0.5):
    theta = np.deg2rad(angle)
    return np.exp(1j * 2 * np.pi * d * np.outer(np.arange(M) , np.sin(theta) / wavelength))

def generate_mutual_coupling_matrix(M,qu):
    """
    生成一个 MxM 的互耦矩阵
    对角线元素为 1，非对角线元素随着距离增加快速衰减，并且是复数
    """
    C = np.eye(M, dtype=complex)  # 初始化为单位矩阵

    for i in range(M):
        for j in range(i + 1, M):
            distance = abs(i - j)
            real_part = np.exp(-distance / 2)  # 实部快速衰减
            imag_part = np.exp(-distance / 2)  # 虚部快速衰减
            C[i, j] = real_part + 1j * imag_part
            C[j, i] = C[i, j]  # 对称性
    
    # 设置左下角和右上角部分为零
    for i in range(M):
        for j in range(M):
            if (j >= M - qu and i <= j - (M - qu)) or (i >= M - qu and j <= i - (M - qu)):
                C[i, j] = 0
    
    return C

def generate_complex_matrix(M):
    """
    生成 MxM 的矩阵 C
    对角线元素为 1，其他非对角线元素为复数，并且所有斜对角线元素相同
    """
    # 初始化矩阵 C 为 MxM 的单位矩阵
    C = np.eye(M, dtype=complex)

    # 生成复数元素
    for i in range(1, M):
        real_part = np.random.rand()  # 随机实部
        imag_part = np.random.rand()  # 随机虚部
        complex_value = real_part + 1j * imag_part  # 生成复数

        for j in range(M - i):
            C[j, j + i] = complex_value  # 上三角斜对角线
            C[j + i, j] = complex_value  # 下三角斜对角线
    return C

def awgn(signal, snr_db):
    snr_linear = 10 ** (snr_db / 10.0)
    power_signal = np.sum(np.abs(signal) ** 2) / signal.size
    power_noise = power_signal / snr_linear
    noise = np.sqrt(power_noise / 2) * (np.random.randn(*signal.shape) + 1j * np.random.randn(*signal.shape))
    return signal + noise

def generate_mat_J(mat_C, q, M):
    mat_J = np.zeros((M, M))#, #dtype=complex)
    for i in range(M):
        for j in range(M):
            if mat_C[i, j] == q:
                mat_J[i, j] = 1
    return mat_J

def gMat_J_all(M,mat_C,Q):
    Mat_all_J = np.zeros((M, M * Q))
    rows = mat_C.shape[0]
    for q in range(Q):
        mat_J = generate_mat_J(mat_C, q + 1, rows)
        # 将 mat_J 按列堆叠到 Mat_all_J 中
        start_col = q * M
        end_col = (q + 1) * M
        Mat_all_J[:, start_col:end_col] = mat_J
    return Mat_all_J

def gMat_T(M, DoA_i, mat_C,Q_len):
    d = 0.5  # number of target
    Geo_Tx = np.arange(M)  # geometry of Tx
    rows = mat_C.shape[0]
    Q = Q_len
    Mat_T = np.zeros((M, Q), dtype=complex)
    for q in range(Q):
        mat_J = generate_mat_J(mat_C, q + 1, rows)  # MATLAB 1-indexed, Python 0-indexed
        At = np.exp(1j * 2 * np.pi * d * Geo_Tx * np.sin(np.deg2rad(DoA_i)))
        Jq = mat_J @ At
        Mat_T[:, q] = Jq
    return Mat_T

# MC互耦阵列的索引
def gMatC(size):
    mat = np.zeros((size, size), dtype=float)
    last_s = 1
    for i in range(size):
        for j in range(size):
            if i==j:
                mat[i,j]=1
            #elif j<i:
            else:
                mat[i, j] = 1+abs(i - j)
    return mat

def gSig(M,Doa,L,MCM,isMc,K=1,SNR=60):
    a = steering_vector(M,Doa)
    S = np.random.randn(K, L) + 1j * np.random.randn(K, L)  # 信号矩阵    
    if isMc==1:
        sig = MCM @ a @ S # 生成信号
    else:
        sig = a @ S
    X = awgn(sig, SNR)  # 添加噪声后信号
    R = (X @ X.conj().T) / L  # 信号协方差矩阵
    D, En = eigh(R, subset_by_index=[0, M-K-1]) # 噪声子空间与噪声特征值
    # 对协方差矩阵进行完整的特征值分解
    D, E = eigh(R)
    # En2 = E[:, :M-K] # 提取噪声子空间
    Es = E[:, M-K:] # 提取信号子空间
    return En,Es
    
# MUSIC_test(M,L,60)    
def MUSIC_test(M,L,DOA):
    En = gSig(M,DOA,L,np.eye(M,M))
    search_grid = np.arange(-90, 90, 1)  # 搜索范围
    G = np.zeros(len(search_grid))
    for idx,theta in enumerate(search_grid):
        a = steering_vector(M,theta)
        G[idx] = 10*np.log10(1/abs(a.conj().T @ En @ En.conj().T @ a))
    plt.figure()
    plt.plot(search_grid,G)
    plt.show()    

tflag = 0
def Pmu(En_lst,MC_gu,M):
    global tflag
    P = 0 # 越大越好
    Doa_s = []
    search_grid = np.arange(-90, 90, 1)  # 搜索范围
    for En in (En_lst):
        gMax = 0
        tdoa = -999
        g_lst = []
        for idx,theta in enumerate(search_grid):
            a = steering_vector(M,theta)
            g = (1/abs(a.conj().T @ MC_gu.conj().T @ En @ En.conj().T @ MC_gu @ a))
            g_lst.append(g[0])
            if gMax<g:
                gMax = g
                tdoa = theta
        # if tflag==1:
        #     plt.figure()
        #     plt.plot(search_grid,g_lst)
        #     plt.show()
        P += gMax
        Doa_s.append(tdoa)
    
    tflag = 1
        
    return Doa_s,P
   
def updateC2(phi,M,w,mat_C,En_lst,Q_len,Doa_zhi,doa_len):
    Q_sum = 0
    for idx in range(doa_len):
        En = En_lst[idx]
        Mat_T = gMat_T(M, Doa_zhi, mat_C,Q_len)
        Q_mat = Mat_T.conj().T @ (En @ En.conj().T) @ Mat_T
        Q_sum += Q_mat
    c = (np.linalg.inv(Q_sum) @ w) /(w.conj().T @ Q_sum @ w)
    c = c/c[0]
    cost_val2 = abs(c.conj().T @ Q_sum @ c)
    psi = np.kron(c,np.eye(M))
    C = phi @ psi
    
    return C,cost_val2  
    
def updateC(phi,M,w,mat_C,En_lst,doa_Gu_lst,Q_len):
    Q_sum = 0
    for idx,doai in enumerate(doa_Gu_lst):
        En = En_lst[idx]
        Mat_T = gMat_T(M, doai, mat_C,Q_len)
        Q_mat = Mat_T.conj().T @ (En @ En.conj().T) @ Mat_T
        Q_sum += Q_mat
    c = (np.linalg.inv(Q_sum) @ w) /(w.conj().T @ Q_sum @ w)
    c = c/c[0]
    cost_val2 = abs(c.conj().T @ Q_sum @ c) # 越小越好
    psi = np.kron(c,np.eye(M))
    C = phi @ psi
    
    return C,cost_val2

def main():
    M = 10
    lambdaf = 1
    d = 0.5
    L = 300
    qu = 6
    k = M - qu
    MCM = generate_mutual_coupling_matrix(M,qu)
    mcm_df = pd.DataFrame(MCM)
    
    # DOAs = generate_DOAs(100)
    # DOAs = np.full(1, 45)
    DOAs = [50]
    doaLen = len(DOAs)
    En_lst = []
    Es_lst = []
    for doa in DOAs:
        En,Es = gSig(M,doa,L,MCM,1)
        En_lst.append(En)
        Es_lst.append(Es)
    mat_C = gMatC(M) # 初始化 MC 下标矩阵
    Q_len = M-qu
    w = np.eye(Q_len, 1)
    phi = gMat_J_all(M,mat_C,Q_len)
    # 初始化MCM
    # 开始设置J1和J2
    Fr = np.hstack((np.zeros((M-2*k,k)),np.eye(M-2*k),np.zeros((M-2*k,k))))
    print("Fr Shape:",Fr.shape)
    #print(Fr)
    nd = 1
    J1 = np.hstack((np.eye(M-nd), np.zeros((M-nd, nd))))
    J2 = np.hstack((np.zeros((M-nd,nd)),np.eye(M-nd)))
    Es = Es_lst[0]
    print("Es Shape:",Es.shape)
    print("J1 Shape:",J1.shape)
    print(Fr@Es)
    Er = Fr@Es
    Er1 = Er[:-1]
    Er2 = Er[1:]
    print("----")
    print(Er1)
    print("-----")
    print(Er2)
    Us1 = J1 @ Es
    Us2 = J2 @ Es
    uu = np.linalg.pinv(Er1) @ Er2
    D, E = np.linalg.eig(uu) #
    
    gu_Doa = np.arcsin(np.angle(D)/(2*np.pi*nd*d/lambdaf))*180/np.pi
    print(gu_Doa)
    #print(np.zeros((M-2, 2)))
    # 还原MCM
    e = np.eye(Q_len, 1)  # d 约束条件
    Mat_T = gMat_T(M, gu_Doa, mat_C,Q_len)
    Q_mat = Mat_T.conj().T @ (En @ En.conj().T) @ Mat_T
    c1 = (np.linalg.inv(Q_mat) @ e)/(e.conj().T @ np.linalg.inv(Q_mat) @ e)
    print(mcm_df.iloc[:][0])
    print(c1)


if __name__=="__main__":
    main()